from flask import Flask, Response, render_template, request, jsonify, send_from_directory
try:
    import paho.mqtt.client as mqtt
except ImportError as e:
    print("❌ MQTT库导入失败，请安装paho-mqtt: pip install paho-mqtt")
    raise e

try:
    import cv2
except ImportError as e:
    print("❌ OpenCV库导入失败，请安装opencv-python或opencv-python-headless:")
    print("   pip install opencv-python-headless")
    raise e

import base64
import numpy as np
import threading
import time
import json
import os

app = Flask(__name__)

# MQTT 配置
broker = 'emqx.digiplc.cn'
port = 1883
control_topic = "camera/control"
feed_topic = "camera/feed"
device_topic = "camera/device"
audio_topic = "camera/audio"  # 新增用于接收音频数据的主题
username = 'test'
password = '!1Qazwsxed'

# 全局变量存储最新图像和设备列表
latest_frames = {}  # 每个设备的最新帧
frame_locks = {}    # 每个设备的锁
devices = {}        # 设备列表
device_lock = threading.Lock()

# 密码配置
CAMERA_PASSWORD = "yyan"

# MQTT 客户端
def on_connect(client, userdata, flags, rc):
    if rc == 0:
        print("✅ 服务器已连接到 MQTT")
        client.subscribe(feed_topic + "/+")
        client.subscribe(device_topic)
    else:
        print(f"❌ MQTT 连接失败: {rc}")

def on_message(client, userdata, msg):
    global latest_frames, frame_locks, devices
    try:
        if msg.topic == device_topic:
            # 快速处理设备状态消息
            try:
                data = json.loads(msg.payload.decode('utf-8'))
                device_id = data.get('device_id')
                status = data.get('status')

                if not device_id:
                    print("⚠️ 忽略无效设备消息: 缺少 device_id")
                    return

                with device_lock:
                    if status == 'online':
                        # 获取推流状态，默认为idle
                        streaming_status = data.get('streaming_status', 'idle')
                        devices[device_id] = {
                            'last_seen': time.time(),
                            'name': data.get('name', f'设备 {device_id}'),
                            'status': streaming_status  # 使用设备发送的推流状态
                        }
                        # 延迟初始化帧存储和锁
                        if device_id not in latest_frames:
                            latest_frames[device_id] = None
                            frame_locks[device_id] = threading.Lock()
                        print(f"📱 设备 {device_id} 上线，状态: {streaming_status}")
                    elif status == 'offline':
                        if device_id in devices:
                            devices.pop(device_id)
                            latest_frames.pop(device_id, None)
                            frame_locks.pop(device_id, None)
                            print(f"📴 设备 {device_id} 下线")
            except json.JSONDecodeError:
                print("❌ 设备消息JSON解析失败")
            except Exception as e:
                print(f"❌ 处理设备消息失败: {e}")

        elif msg.topic.startswith(feed_topic):
            # 提前解析 device_id
            device_id = msg.topic.split('/')[-1]
            if not device_id:
                print("⚠️ 忽略无效图像消息: 无法解析 device_id")
                return

            # 尽量不在锁中执行耗时操作
            try:
                # 解码 base64 图像（客户端已发送 RGB 格式）
                jpg_bytes = base64.b64decode(msg.payload.decode('utf-8'))
                np_arr = np.frombuffer(jpg_bytes, np.uint8)
                img_rgb = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

                if img_rgb is None or img_rgb.size == 0:
                    print(f"❌ 设备 {device_id} 图像解码失败或为空")
                    return

                # 只在最后更新时加锁
                if device_id in frame_locks:
                    with frame_locks[device_id]:
                        latest_frames[device_id] = img_rgb
                else:
                    # 新设备可能未注册，忽略首帧
                    pass

            except (base64.binascii.Error, UnicodeDecodeError):
                print(f"❌ 设备 {device_id} Base64解码失败")
            except Exception as e:
                print(f"❌ 设备 {device_id} 图像处理失败: {e}")

    except Exception as e:
        print(f"❌ 处理MQTT消息失败: {e}")

# 启动 MQTT 客户端
mqtt_client = mqtt.Client()
mqtt_client.username_pw_set(username, password)
mqtt_client.on_connect = on_connect
mqtt_client.on_message = on_message

def start_mqtt():
    mqtt_client.connect(broker, port, 60)
    mqtt_client.loop_forever()

# 启动 MQTT 在后台线程
mqtt_thread = threading.Thread(target=start_mqtt, daemon=True)
mqtt_thread.start()

# Flask 路由
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/verify_password', methods=['POST'])
def verify_password():
    data = request.get_json()
    password = data.get('password') if data else None
    
    if password == CAMERA_PASSWORD:
        return jsonify({'success': True})
    else:
        return jsonify({'success': False}), 401

@app.route('/get_devices')
def get_devices():
    with device_lock:
        # 清理长时间未响应的设备（超过30秒）
        current_time = time.time()
        expired_devices = [
            device_id for device_id, info in devices.items()
            if current_time - info['last_seen'] > 30
        ]
        for device_id in expired_devices:
            devices.pop(device_id, None)
            latest_frames.pop(device_id, None)
            frame_locks.pop(device_id, None)
            print(f"⏰ 设备 {device_id} 超时离线")
        
        return jsonify(devices)

@app.route('/control_camera', methods=['POST'])
def control_camera():
    try:
        data = request.get_json()
        command = data.get('command')
        device_id = data.get('device_id')
        
        if not device_id:
            return jsonify({'success': False, 'error': '未指定设备ID'}), 400
        
        # 发送控制指令到指定摄像头客户端
        control_payload = {
            'command': command
        }
        
        # 如果有设置参数，也一并发送
        if 'settings' in data:
            control_payload['settings'] = data['settings']
            
        # 如果有色彩校正参数，也一并发送
        if 'color_correction' in data:
            control_payload['color_correction'] = data['color_correction']
            
        # 如果是播放音频命令，包含音频文件信息
        if command == 'play_audio':
            audio_info = {
                'audio_file': data.get('audio_file', 'audio/yu.MP3'),
                'volume': data.get('volume', 1.0)
            }
            control_payload['audio'] = audio_info
        
        mqtt_client.publish(f"{control_topic}/{device_id}", json.dumps(control_payload))
        
        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

@app.route('/upload_audio', methods=['POST'])
def upload_audio():
    try:
        if 'audio' not in request.files:
            return jsonify({'success': False, 'error': '没有音频文件'}), 400
        
        file = request.files['audio']
        if file.filename == '':
            return jsonify({'success': False, 'error': '没有选择文件'}), 400
        
        # 保存到audio文件夹
        audio_folder = 'audio'
        if not os.path.exists(audio_folder):
            os.makedirs(audio_folder)
        
        # 修复文件扩展名问题，使用webm格式
        filepath = os.path.join(audio_folder, 'recorded_audio.webm')
        file.save(filepath)
        
        print(f"✅ 音频文件已保存: {filepath}")
        return jsonify({'success': True})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

@app.route('/icons/<path:filename>')
def icons(filename):
    return send_from_directory('templates/icons', filename)

def gen_frames(device_id):
    while True:
        if device_id in frame_locks:
            with frame_locks[device_id]:
                if device_id in latest_frames and latest_frames[device_id] is not None:
                    # 使用更低的JPEG质量
                    encode_params = [cv2.IMWRITE_JPEG_QUALITY, 30]
                    _, buffer = cv2.imencode('.jpg', latest_frames[device_id], encode_params)
                    frame = buffer.tobytes()
                    yield (b'--frame\r\n'
                           b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
        time.sleep(0.067)  # 提高刷新率

@app.route('/stream/<device_id>')
def stream(device_id):
    return Response(gen_frames(device_id),
                    mimetype='multipart/x-mixed-replace; boundary=frame')

# 为Gunicorn添加应用入口点
application = app

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=31989, threaded=True)